import pandas as pd
import numpy as np
import scipy


def basis_1(base_df, comp_df):
  base = base_df[[
      'global_timestamp',
      't.vwap.1m',
      't.vwap.5m',
      't.vwap.10m',
      't.vwap.20m',
      't.vwap.30m',
      'Y.mid_ret.5m'
  ]]
  comp = comp_df[[
      'global_timestamp', 't.vwap.1m', 't.vwap.5m', 't.vwap.10m', 't.vwap.20m', 't.vwap.30m'
  ]]
  df = pd.merge(base, comp, how='inner', on='global_timestamp').dropna()
  print(df)
  # x = (df['t.vwap.20m_y'] / df['t.vwap.20m_x']) / (df['t.vwap.5m_y'] / df['t.vwap.5m_x'])
  x = (df['t.vwap.20m_x'] - df['t.vwap.5m_x']) + (df['t.vwap.20m_y'] - df['t.vwap.5m_y'])

  x[np.isnan(x)] = 0
  x[(x > 100) | (x < -100)] = 0
  print('------------------')
  idx = x.index
  x = scipy.stats.zscore(x)
  x = pd.Series(x, index=idx)

  x[(x > 1)] = 1
  x[(x < -1)] = -1

  # x = np.cbrt(x)

  # x = (df['t.vwap.30m_y'].shift(1) - df['t.vwap.30m_y'])
  y = df['Y.mid_ret.5m']
  print(y)
  t = df['global_timestamp']
  return x, y, t


def basis_2(base_df, comp_df):
  base = base_df[[
      'global_timestamp',
      't.vwap.1m',
      't.vwap.5m',
      't.vwap.10m',
      't.vwap.20m',
      't.vwap.30m',
      'Y.mid_ret.5m'
  ]]
  comp = comp_df[[
      'global_timestamp', 't.vwap.1m', 't.vwap.5m', 't.vwap.10m', 't.vwap.20m', 't.vwap.30m'
  ]]
  df = pd.merge(base, comp, how='inner', on='global_timestamp').dropna()
  print(df)
  x = (df['t.vwap.10m_y'] - df['t.vwap.1m_y']) + (df['t.vwap.30m_y'] - df['t.vwap.20m_y'])

  x[np.isnan(x)] = 0
  x[(x > 100) | (x < -100)] = 0
  idx = x.index
  x = scipy.stats.zscore(x)
  x = pd.Series(x, index=idx)
  x[(x > 1)] = 1
  x[(x < -1)] = -1
  y = df['Y.mid_ret.5m']
  t = df['global_timestamp']
  return x, y, t


def basis_3(base_df, comp_df):
  base = base_df[[
      'global_timestamp',
      't.vwap.1m',
      't.vwap.5m',
      't.vwap.10m',
      't.vwap.20m',
      't.vwap.30m',
      'Y.mid_ret.5m'
  ]]
  comp = comp_df[[
      'global_timestamp', 't.vwap.1m', 't.vwap.5m', 't.vwap.10m', 't.vwap.20m', 't.vwap.30m'
  ]]
  df = pd.merge(base, comp, how='inner', on='global_timestamp').dropna()
  print(df)
  x = (df['t.vwap.20m_y'] - df['t.vwap.5m_y'])

  x[np.isnan(x)] = 0
  x[(x > 100) | (x < -100)] = 0
  idx = x.index
  x = scipy.stats.zscore(x)
  x = pd.Series(x, index=idx)
  x[(x > 1)] = 1
  x[(x < -1)] = -1
  y = df['Y.mid_ret.5m']
  t = df['global_timestamp']
  return x, y, t
